import cv2
import numpy as np
import time
import math

# 初始化摄像头
cap = cv2.VideoCapture(4)
if not cap.isOpened():
    print("无法打开摄像头！")
    exit()

# 设置摄像头参数
cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)
cap.set(cv2.CAP_PROP_FPS, 30)

# 相机参数（根据实际相机校准）
fx = 800  # 焦距 (x方向)
fy = 800  # 焦距 (y方向)
cx = 320  # 主点 (x方向)
cy = 240  # 主点 (y方向)

# 实际物体尺寸（单位：米） - 需要根据实际物体测量
REAL_BALL_DIAMETER = 0.1    # 黄色球体实际直径
REAL_CUBE_HEIGHT = 0.15     # 绿色立方体实际高度
REAL_CYLINDER_HEIGHT = 0.2  # 绿色圆柱体实际高度

# 创建窗口
cv2.namedWindow('Multi-Object Detection', cv2.WINDOW_NORMAL)
cv2.namedWindow('Color Mask', cv2.WINDOW_NORMAL)

# 优化后的颜色范围定义（HSV空间）
lower_green = np.array([35, 70, 70])    # 绿色物体
upper_green = np.array([85, 255, 255])
lower_yellow = np.array([20, 120, 120]) # 黄色物体
upper_yellow = np.array([30, 255, 255])

# 检测参数优化
min_contour_area = 200          # 降低最小轮廓面积以适应遮挡
ball_min_circularity = 0.65     # 降低圆度要求以适应遮挡
cylinder_min_aspect_ratio = 0.5 # 放宽宽高比范围

# 形态学核
kernel_small = np.ones((3,3), np.uint8)
kernel_medium = np.ones((5,5), np.uint8)

def calculate_depth(obj_type, size_param):
    """
    根据物体类型和尺寸参数计算深度
    :param obj_type: 物体类型 ('ball', 'cube', 'cylinder')
    :param size_param: 对于球体是半径，对于立方体/圆柱体是高度
    :return: 深度值(米)
    """
    if obj_type == 'ball':
        # 基于球体直径计算深度
        if size_param > 0:
            Z = (fx * REAL_BALL_DIAMETER) / (2 * size_param)
            return max(0.1, min(5.0, Z))  # 限制在0.1-5米范围内
    elif obj_type == 'cube':
        # 基于立方体高度计算深度
        if size_param > 0:
            Z = (fy * REAL_CUBE_HEIGHT) / size_param
            return max(0.1, min(5.0, Z))
    elif obj_type == 'cylinder':
        # 基于圆柱体高度计算深度
        if size_param > 0:
            Z = (fy * REAL_CYLINDER_HEIGHT) / size_param
            return max(0.1, min(5.0, Z))
    return 1.0  # 默认值

def detect_objects(hsv_frame, color_frame):
    # 创建颜色掩膜
    green_mask = cv2.inRange(hsv_frame, lower_green, upper_green)
    yellow_mask = cv2.inRange(hsv_frame, lower_yellow, upper_yellow)
    
    # 优化形态学操作
    green_mask = cv2.morphologyEx(green_mask, cv2.MORPH_OPEN, kernel_small)
    green_mask = cv2.morphologyEx(green_mask, cv2.MORPH_CLOSE, kernel_medium)
    yellow_mask = cv2.morphologyEx(yellow_mask, cv2.MORPH_OPEN, kernel_small)
    yellow_mask = cv2.morphologyEx(yellow_mask, cv2.MORPH_CLOSE, kernel_medium)
    
    # 组合掩膜用于显示
    color_mask = cv2.bitwise_or(green_mask, yellow_mask)
    
    detected_objects = []
    
    # 检测绿色物体（圆柱/方块）
    green_contours, _ = cv2.findContours(green_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    for contour in green_contours:
        area = cv2.contourArea(contour)
        if area < min_contour_area:
            continue
            
        # 计算轮廓特征
        perimeter = cv2.arcLength(contour, True)
        circularity = 4 * np.pi * area / (perimeter**2) if perimeter > 0 else 0
        x, y, w, h = cv2.boundingRect(contour)
        aspect_ratio = float(w)/h
        
        # 计算中心点
        center_x = x + w//2
        center_y = y + h//2
        
        # 判断物体类型
        if circularity > ball_min_circularity and 0.7 < aspect_ratio < 1.4:
            rect = cv2.minAreaRect(contour)
            box = cv2.boxPoints(rect)
            box = np.int0(box)
            
            side1 = np.linalg.norm(box[0]-box[1])
            side2 = np.linalg.norm(box[1]-box[2])
            long_side = max(side1, side2)
            short_side = min(side1, side2)
            aspect_ratio = long_side / short_side if short_side > 0 else 0
            
            if aspect_ratio > 1.5:
                obj_type = "cylinder"
                Z = calculate_depth(obj_type, h)
            else:
                obj_type = "cube"
                Z = calculate_depth(obj_type, h)
            
            # 计算3D坐标
            X = (center_x - cx) * Z / fx
            Y = (center_y - cy) * Z / fy
            
            detected_objects.append({
                'type': obj_type,
                'contour': contour,
                'bbox': (x, y, w, h),
                'center': (center_x, center_y),
                'world_coords': (X, Y, Z),
                'color': 'green'
            })
    
    # 检测黄色球体
    yellow_contours, _ = cv2.findContours(yellow_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    for contour in yellow_contours:
        area = cv2.contourArea(contour)
        if area < min_contour_area:
            continue
            
        hull = cv2.convexHull(contour)
        hull_area = cv2.contourArea(hull)
        solidity = float(area)/hull_area if hull_area > 0 else 0
        
        (x, y), radius = cv2.minEnclosingCircle(contour)
        center = (int(x), int(y))
        radius = int(radius)
        
        if solidity > 0.7 or (radius > 10 and area > 0.5 * math.pi * radius**2):
            Z = calculate_depth('ball', radius)
            
            X = (center[0] - cx) * Z / fx
            Y = (center[1] - cy) * Z / fy
            
            detected_objects.append({
                'type': 'ball',
                'center': center,
                'radius': radius,
                'world_coords': (X, Y, Z),
                'color': 'yellow'
            })
    
    return detected_objects, color_mask

# 主循环
frame_count = 0
start_time = time.time()

while True:
    ret, color_frame = cap.read()
    if not ret:
        print("无法获取画面")
        break
    
    hsv_frame = cv2.cvtColor(color_frame, cv2.COLOR_BGR2HSV)
    detected_objects, color_mask = detect_objects(hsv_frame, color_frame)
    
    # 输出物体坐标
    for obj in detected_objects:
        X, Y, Z = obj['world_coords']
        print(f"{obj['type']}: X={X:.3f}m, Y={Y:.3f}m, Z={Z:.3f}m")
    
    # 绘制结果
    display_frame = color_frame.copy()
    for obj in detected_objects:
        obj_type = obj['type']
        X, Y, Z = obj['world_coords']
        
        if obj_type == 'ball':
            center = obj['center']
            radius = obj['radius']
            cv2.circle(display_frame, center, radius, (0, 255, 255), 2)
            cv2.drawMarker(display_frame, center, (0, 0, 255), 
                         markerType=cv2.MARKER_CROSS, markerSize=20, thickness=2)
            cv2.putText(display_frame, f"Ball: ({X:.2f}, {Y:.2f}, {Z:.2f})", 
                       (center[0]-30, center[1]-radius-10), 
                       cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
        elif obj_type == 'cube':
            x, y, w, h = obj['bbox']
            cv2.rectangle(display_frame, (x, y), (x+w, y+h), (0, 255, 0), 2)
            cv2.putText(display_frame, f"Cube: ({X:.2f}, {Y:.2f}, {Z:.2f})", 
                       (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
        elif obj_type == 'cylinder':
            x, y, w, h = obj['bbox']
            center = obj['center']
            cv2.ellipse(display_frame, (center[0], center[1]), (w//2, h//2), 
                       0, 0, 360, (0, 255, 0), 2)
            cv2.putText(display_frame, f"Cylinder: ({X:.2f}, {Y:.2f}, {Z:.2f})", 
                       (x, y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
        
        # 绘制中心点
        if 'center' in obj:
            cv2.drawMarker(display_frame, obj['center'], (0, 0, 255),
                          markerType=cv2.MARKER_CROSS, markerSize=20, thickness=2)
    
    # 显示帧率
    frame_count += 1
    if frame_count % 10 == 0:
        fps = 10 / (time.time() - start_time)
        start_time = time.time()
        cv2.putText(display_frame, f"FPS: {fps:.1f}", (10, 30), 
                   cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255), 2)
    
    cv2.imshow('Multi-Object Detection', display_frame)
    cv2.imshow('Color Mask', color_mask)
    
    if cv2.waitKey(1) == ord('q'):
        break

cap.release()
cv2.destroyAllWindows()
